import time
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
from gensim.models import Word2Vec
from tqdm import tqdm
import psutil
import os
import gc
from faker import Faker
import jieba

class AddressVectorization:
    def __init__(self, model_type='sbert', model_name='paraphrase-mpnet-base-v2', batch_size=256):
        """初始化向量化处理器"""
        self.model_type = model_type
        self.model_name = model_name
        self.batch_size = batch_size
        self.faker = Faker(['zh_CN'])
        
        if model_type == 'sbert':
            self.model = SentenceTransformer(model_name)
        elif model_type == 'word2vec':
            # 初始化Word2Vec模型
            self.model = None  # 将在处理数据时训练
        else:
            raise ValueError(f"不支持的模型类型: {model_type}")
        
    def get_memory_usage(self):
        """获取当前内存使用情况（MB）"""
        gc.collect()
        process = psutil.Process(os.getpid())
        return process.memory_info().rss / 1024 / 1024

    def generate_realistic_addresses(self, n):
        """生成更真实的中国省市区地址数据"""
        # 定义常用省份和对应的主要城市
        provinces_cities = {
            '北京市': ['朝阳区', '海淀区', '丰台区', '西城区', '东城区', '昌平区', '通州区'],
            '上海市': ['浦东新区', '徐汇区', '长宁区', '静安区', '普陀区', '虹口区', '杨浦区'],
            '广东省': {
                '广州市': ['天河区', '越秀区', '海珠区', '白云区', '黄埔区'],
                '深圳市': ['南山区', '福田区', '罗湖区', '宝安区', '龙岗区']
            },
            '浙江省': {
                '杭州市': ['西湖区', '上城区', '下城区', '江干区', '拱墅区'],
                '宁波市': ['海曙区', '江北区', '北仑区', '镇海区', '鄞州区']
            },
            '江苏省': {
                '南京市': ['玄武区', '秦淮区', '建邺区', '鼓楼区', '浦口区'],
                '苏州市': ['姑苏区', '虎丘区', '吴中区', '相城区', '工业园区']
            }
        }
        
        # 定义常用路名和建筑物
        street_names = ['人民路', '中山路', '解放路', '建设路', '和平路', '新华路', '长江路', '东风路']
        building_types = ['大厦', '广场', '中心', '大楼', '花园', '公寓', '小区']
        
        addresses = []
        for _ in tqdm(range(n), desc=f"生成{n}条地址"):
            # 随机选择省份
            province = np.random.choice(list(provinces_cities.keys()))
            
            if isinstance(provinces_cities[province], list):
                # 直辖市
                city = province
                district = np.random.choice(provinces_cities[province])
            else:
                # 普通省份
                city = np.random.choice(list(provinces_cities[province].keys()))
                district = np.random.choice(provinces_cities[province][city])
            
            # 生成详细地址
            street = np.random.choice(street_names)
            building = np.random.choice(building_types)
            building_number = np.random.randint(1, 999)
            detail_number = np.random.randint(1, 50)
            
            # 组合完整地址
            address_parts = [
                province,
                city if city != province else '',
                district,
                f"{street}{building_number}号",
                f"{np.random.choice(['金色', '银座', '环球', '东方', '国际'])}",
                f"{building}",
                f"{detail_number}号"
            ]
            
            # 过滤空字符串并拼接
            address = ''.join(filter(None, address_parts))
            addresses.append(address)
        
        return addresses

    def process_batch(self, addresses):
        """批量处理地址向量化"""
        if self.model_type == 'sbert':
            return self.model.encode(
                addresses,
                batch_size=self.batch_size,
                show_progress_bar=True,
                convert_to_numpy=True
            )
        elif self.model_type == 'word2vec':
            # 对地址进行分词
            segmented_addresses = [list(jieba.cut(addr)) for addr in addresses]
            
            # 训练Word2Vec模型
            if self.model is None:
                self.model = Word2Vec(
                    sentences=segmented_addresses,
                    vector_size=300,
                    window=5,
                    min_count=1,
                    workers=4
                )
            
            # 计算每个地址的平均词向量
            vectors = []
            for words in segmented_addresses:
                word_vectors = [self.model.wv[word] for word in words if word in self.model.wv]
                if word_vectors:
                    address_vector = np.mean(word_vectors, axis=0)
                else:
                    address_vector = np.zeros(self.model.vector_size)
                vectors.append(address_vector)
            
            return np.array(vectors)

    def run_benchmark(self, sizes=[100, 1000]):
        """运行性能测试"""
        results = []
        
        for size in sizes:
            print(f"\n{'='*50}")
            print(f"测试数据量: {size:,} | 模型: {self.model_type}-{self.model_name}")
            
            initial_memory = self.get_memory_usage()
            
            try:
                # 生成数据并计时
                gen_start = time.time()
                addresses = self.generate_realistic_addresses(size)
                gen_time = time.time() - gen_start
                
                after_gen_memory = self.get_memory_usage()
                
                # 向量化处理
                vec_start = time.time()
                vectors = self.process_batch(addresses)
                vec_time = time.time() - vec_start
                
                total_memory = self.get_memory_usage() - initial_memory
                
                result = {
                    "模型": f"{self.model_type}-{self.model_name}",
                    "数据量": f"{size:,}",
                    "数据生成时间(秒)": round(gen_time, 2),
                    "向量化时间(秒)": round(vec_time, 2),
                    "总耗时(秒)": round(gen_time + vec_time, 2),
                    "每秒处理条数": round(size / (gen_time + vec_time), 2),
                    "数据生成内存(MB)": round(after_gen_memory - initial_memory, 2),
                    "总内存占用(MB)": round(total_memory, 2),
                    "向量维度": vectors.shape[1]
                }
                
                results.append(result)
                print("\n当前批次结果:")
                for key, value in result.items():
                    print(f"{key}: {value}")
                
            except Exception as e:
                print(f"处理{size}条数据时出错: {str(e)}")
                
            finally:
                if 'addresses' in locals(): del addresses
                if 'vectors' in locals(): del vectors
                gc.collect()
        
        return pd.DataFrame(results)

def format_table(df, is_summary=False):
    """美化表格展示"""
    # 设置样式
    styler = df.style
    
    # 基础数字格式化
    format_dict = {
        '向量化时间(秒)': '{:.2f}',
        '处理速度(条/秒)': '{:.1f}',
        '内存占用(MB)': '{:.1f}',
        '向量维度': '{:,.0f}'
    }
    
    if is_summary:
        format_dict = {
            '平均向量化时间(秒)': '{:.2f}',
            '时间标准差': '{:.2f}',
            '平均处理速度(条/秒)': '{:.1f}',
            '速度标准差': '{:.1f}',
            '平均内存占用(MB)': '{:.1f}',
            '峰值内存(MB)': '{:.1f}'
        }
    
    styler = styler.format(format_dict)
    
    # 设置背景色突出显示
    if is_summary:
        highlight_cols = {
            'min': ['平均向量化时间(秒)', '平均内存占用(MB)'],
            'max': ['平均处理速度(条/秒)']
        }
    else:
        highlight_cols = {
            'min': ['向量化时间(秒)', '内存占用(MB)'],
            'max': ['处理速度(条/秒)']
        }
    
    # 添加背景色
    for method, cols in highlight_cols.items():
        if method == 'min':
            styler = styler.highlight_min(subset=cols, color='#90EE90', axis=0)  # 浅绿色
        else:
            styler = styler.highlight_max(subset=cols, color='#90EE90', axis=0)
    
    # 添加表格边框
    styler = styler.set_table_styles([
        {'selector': 'th', 'props': [('background-color', '#f0f0f0'),
                                   ('font-weight', 'bold'),
                                   ('text-align', 'center'),
                                   ('border', '1px solid #ddd'),
                                   ('padding', '8px')]},
        {'selector': 'td', 'props': [('border', '1px solid #ddd'),
                                   ('padding', '8px'),
                                   ('text-align', 'right')]},
        {'selector': '', 'props': [('border-collapse', 'collapse'),
                                 ('font-family', 'Arial, sans-serif')]}
    ])
    
    return styler

def compare_models():
    """比较不同模型的性能"""
    test_sizes = [100, 1000, 5000, 10000]  # 可根据需要调整测试规模
    
    model_configs = [
        ('sbert', 'paraphrase-mpnet-base-v2'),          # SBERT模型
        ('sbert', 'BAAI/bge-base-zh-v1.5'),                 # BGE中文模型
        ('word2vec', 'word2vec-chinese'),               # Word2Vec模型
    ]
    
    all_results = []
    
    try:
        for model_type, model_name in model_configs:
            print(f"\n开始测试模型: {model_type}-{model_name}")
            processor = AddressVectorization(
                model_type=model_type,
                model_name=model_name,
                batch_size=256
            )
            results_df = processor.run_benchmark(test_sizes)
            all_results.append(results_df)
        
        # 合并所有结果
        final_results = pd.concat(all_results, ignore_index=True)
        
        # 保存原始结果
        final_results.to_csv('model_comparison_results.csv', index=False)
        
        # 创建性能对比表格
        print("\n============ 模型性能对比分析 ============")
        
        # 1. 按数据量分组的性能对比
        for size in test_sizes:
            size_str = f"{size:,}"
            print(f"\n数据量: {size_str} 条")
            
            # 获取当前数据量的结果
            size_results = final_results[final_results['数据量'] == size_str]
            
            # 创建对比表格
            comparison_metrics = {
                '向量化时间(秒)': '向量化时间(秒)',
                '处理速度(条/秒)': '每秒处理条数',
                '内存占用(MB)': '总内存占用(MB)',
                '向量维度': '向量维度'
            }
            
            # 创建更美观的表格
            comparison_table = pd.DataFrame()
            
            for metric_name, column in comparison_metrics.items():
                row = size_results.set_index('模型')[column]
                comparison_table[metric_name] = row
            
            # 打印对比表格
            print("\n性能指标对比:")
            styled_table = format_table(comparison_table)
            print(styled_table.to_string())
            
            # 计算相对性能并美化展示
            print("\n相对性能比较（基准：BGE模型）:")
            bge_model = 'sbert-BAAI/bge-base-zh-v1.5'
            relative_performance = comparison_table.div(comparison_table.loc[bge_model])
            
            # 添加性能说明列
            performance_with_desc = pd.DataFrame()
            for col in relative_performance.columns:
                if col != '向量维度':
                    # 数值列
                    performance_with_desc[col] = relative_performance[col].round(2)
                    # 说明列
                    performance_with_desc[f'{col}\n(相对表现)'] = relative_performance[col].apply(
                        lambda x: '✓ 更好' if x < 1 else ('✗ 更差' if x > 1 else '━ 相同')
                    )
            
            # 美化相对性能表格
            styled_relative = performance_with_desc.style.format({
                col: '{:.2f}x' for col in relative_performance.columns if col != '向量维度'
            })
            print(styled_relative.to_string())
            
        # 汇总统计表格美化
        print("\n============ 汇总统计 ============")
        summary_stats = final_results.groupby('模型').agg({
            '向量化时间(秒)': ['mean', 'std'],
            '每秒处理条数': ['mean', 'std'],
            '总内存占用(MB)': ['mean', 'max']
        }).round(2)
        
        summary_stats.columns = [
            '平均向量化时间(秒)', '时间标准差',
            '平均处理速度(条/秒)', '速度标准差',
            '平均内存占用(MB)', '峰值内存(MB)'
        ]
        
        styled_summary = format_table(summary_stats, is_summary=True)
        print("\n模型整体性能统计:")
        print(styled_summary.to_string())
        
        # 3. 生成可视化图表
        try:
            import matplotlib.pyplot as plt
            import seaborn as sns
            
            # 创建模型名称映射
            model_name_map = {
                'sbert-paraphrase-mpnet-base-v2': 'SBERT',
                'sbert-BAAI/bge-base-zh-v1.5': 'BGE-base',
                'word2vec-word2vec-chinese': 'Word2Vec'
            }
            
            # 复制数据并替换模型名称
            plot_data = final_results.copy()
            plot_data['Model'] = plot_data['模型'].map(model_name_map)
            
            # 设置图表风格
            sns.set_style("whitegrid")
            plt.figure(figsize=(15, 10))
            
            # 处理速度对比
            plt.subplot(2, 2, 1)
            ax1 = sns.barplot(data=plot_data, x='Model', y='每秒处理条数')
            plt.title('Processing Speed Comparison', fontsize=12, pad=15)
            plt.xlabel('Model')
            plt.ylabel('Records per Second')
            plt.xticks(rotation=30, ha='right')
            for i in ax1.containers:
                ax1.bar_label(i, fmt='%.1f', padding=3)
            
            # 内存占用对比
            plt.subplot(2, 2, 2)
            ax2 = sns.barplot(data=plot_data, x='Model', y='总内存占用(MB)')
            plt.title('Memory Usage Comparison', fontsize=12, pad=15)
            plt.xlabel('Model')
            plt.ylabel('Memory Usage (MB)')
            plt.xticks(rotation=30, ha='right')
            for i in ax2.containers:
                ax2.bar_label(i, fmt='%.1f', padding=3)
            
            # 向量维度对比
            plt.subplot(2, 2, 3)
            ax3 = sns.barplot(data=plot_data, x='Model', y='向量维度')
            plt.title('Vector Dimension Comparison', fontsize=12, pad=15)
            plt.xlabel('Model')
            plt.ylabel('Dimension')
            plt.xticks(rotation=30, ha='right')
            for i in ax3.containers:
                ax3.bar_label(i, fmt='%.0f', padding=3)
            
            # 调整布局
            plt.tight_layout()
            
            # 保存图表
            plt.savefig('model_comparison_charts.png', 
                       bbox_inches='tight', 
                       dpi=300,
                       facecolor='white')
            print("\n性能对比图表已保存至 model_comparison_charts.png")
            
        except ImportError:
            print("\n注: 如需生成可视化图表，请安装 matplotlib 和 seaborn")
        
        print("\n详细结果已保存至 model_comparison_results.csv")
        
    except Exception as e:
        print(f"测试过程中发生错误: {str(e)}")

if __name__ == "__main__":
    compare_models()